import networkx as nx
import matplotlib.pyplot as plt
from enthought.mayavi import mlab
def drawG(nodelist, edgelist, loc, filename):
    G = nx.Graph()
    G.add_nodes_from(nodelist)
    G.add_edges_from(edgelist)
    nx.draw(G, loc)
    plt.savefig(filename)
    plt.clf()

import numpy as np
def draw3DPickerGraph(nodelist, edgelist, pos):    
    global xyz, G, H, scalars
    
    xyz=pos
# scalar colors
    scalars=np.array(nodelist)+5

# Setup the Env
    figure = mlab.figure(1, bgcolor=(0, 0, 0))
    mlab.clf()
    figure.scene.disable_render = True
    
    pts = mlab.points3d([i[0] for i in xyz],
                        [i[1] for i in xyz],
                        [i[2] for i in xyz],
                        scalars,
                        scale_factor=0.1,
                        scale_mode='none',
                        colormap='Blues',
                        resolution=20)

    pts.mlab_source.dataset.lines = np.array(edgelist)
    tube = mlab.pipeline.tube(pts, tube_radius=0.01)
    mlab.pipeline.surface(tube, color=(0.8, 0.8, 0.8))
    
# Add an outline to show the selected point and center it on the first
# data point.
    outline = mlab.outline(line_width=3)
    outline.outline_mode = 'cornered'
    outline.bounds = (pos[0][0]-0.1, pos[0][0]+0.1,
                      pos[0][1]-0.1, pos[0][1]+0.1,
                      pos[0][2]-0.1, pos[0][2]+0.1)    
    a = """
    outline.bounds = (pos[0][0]-0.1, pos[0][0]+0.1,
                      pos[0][1]-0.1, pos[0][1]+0.1,
                      pos[0][2]-0.1, pos[0][2]+0.1)
    """
# Every object has been created, we can reenable the rendering.
    figure.scene.disable_render = False
################################################################################

# Here, we grab the points describing the individual glyph, to figure
# out how many points are in an individual glyph.
    glyph_points = pts.glyph.glyph_source.glyph_source.output.points.to_array()

    def picker_callback(picker):
        """ Picker callback: this get called when on pick events.
        """
        if picker.actor in pts.actor.actors:
            # Find which data point corresponds to the point picked:
            # we have to account for the fact that each data point is
            # represented by a glyph with several points
            point_id = picker.point_id/glyph_points.shape[0]
            # If the no points have been selected, we have '-1'
            if point_id != -1:
                # Retrieve the coordinnates coorresponding to that data
                # point
                x, y, z = pos[point_id]
                # Move the outline to the data point.
                outline.bounds = (x-0.1, x+0.1,
                                  y-0.1, y+0.1,
                                  z-0.1, z+0.1)


    picker = figure.on_mouse_pick(picker_callback)

# Decrease the tolerance, so that we can more easily select a precise
# point.
    picker.tolerance = 0.01
    #mlab.title('Click on red balls')    
    mlab.show()

    
    
